580C - Kefa and Park - CodeForces Solution


dfs and similar graphs trees *1500

Please click on ads to support us..

Python Code:

'''
Don't Copy This Code, CopyRight . [email protected] © 2022-2023 :)
'''

import sys
import threading
sys.setrecursionlimit(2147483647)
input = sys.stdin.readline
def print(*args, end='\n', sep=' ') -> None:
    sys.stdout.write(sep.join(map(str, args)) + end)

def dfs(node , visited , c):
    global graph, x
    if(x[node - 1]):c += 1
    else:c = 0
    
    if(c > m):return 0
    
    ans = 0
    for i in graph[node]:
 
        if(i != visited):
            ans += dfs(i , node , c)
 
    if(node != 1 and len(graph[node]) == 1):return 1
    return ans

def Solve():
    global x, m, ans, graph
    n, m = map(int, input().split())
    x = list(map(int, input().split()))
    ans = 0
    graph = dict()
    for i in range(n-1):
        a, b = map(int, input().split())
        if a not in graph.keys():
            graph[a] = set()
        if b not in graph.keys():
            graph[b] = set()
        graph[a].add(b)
        graph[b].add(a)

    ans = dfs(1, -1, 0)
    print(ans)


if __name__ == "__main__":
    threading.stack_size(10**8)
    threading.Thread(target=Solve).start()

C++ Code:

#include <bits/stdc++.h>
#define pb push_back
#define sz(x) ((int)(x).size())
#define all(x) (x).begin(), (x).end()
#define ll long long
using namespace std;
const int MOD = 1e9 + 7;
const int N = 2e5+6;

int n,m;
vector<int>g[N];
int a[N];

int dfs(int u,int c,int p){
    if(c > m)return 0;
    if(sz(g[u]) == 1 && u != 1)return 1;
    int ret = 0;
    for(auto v : g[u]){
        if(v == p)continue;
        if(a[v] == 1){
            ret += dfs(v,c+1,u);
        }else{
            ret += dfs(v,0,u);
        }
    }
    return ret;
}
void solve(int tc){
    scanf("%d %d",&n,&m);
    for(int i = 1;i<=n;i++)scanf("%d",&a[i]);
    for(int i = 1,u,v;i<n;i++)scanf("%d %d",&u,&v),g[u].pb(v),g[v].pb(u);
    printf("%d\n",dfs(1,a[1],0));
}

int main(){
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    int tc = 1;
    //cin>>tc;
    //scanf("%d",&tc);
    for(int i = 1;i<=tc;i++)
        solve(i);
}


Comments

Submit
0 Comments
More Questions

1613B - Absent Remainder
1536B - Prinzessin der Verurteilung
1699B - Almost Ternary Matrix
1545A - AquaMoon and Strange Sort
538B - Quasi Binary
424A - Squats
1703A - YES or YES
494A - Treasure
48B - Land Lot
835A - Key races
1622C - Set or Decrease
1682A - Palindromic Indices
903C - Boxes Packing
887A - Div 64
755B - PolandBall and Game
808B - Average Sleep Time
1515E - Phoenix and Computers
1552B - Running for Gold
994A - Fingerprints
1221C - Perfect Team
1709C - Recover an RBS
378A - Playing with Dice
248B - Chilly Willy
1709B - Also Try Minecraft
1418A - Buying Torches
131C - The World is a Theatre
1696A - NIT orz
1178D - Prime Graph
1711D - Rain
534A - Exam